package com.xstudio.pvzreborn.world;

import com.mojang.datafixers.util.Pair;
import com.xstudio.pvzreborn.PVZReborn;
import com.xstudio.pvzreborn.world.structure.KnightTomb;
import com.xstudio.pvzreborn.world.structure.NetherTower;
import com.xstudio.pvzreborn.world.structure.StructureFeatureBase;
import com.xstudio.pvzreborn.world.tree.TreeFeatures;
import com.xstudio.pvzreborn.world.tree.TreeFeatures.Placements;
import net.minecraft.core.Holder;
import net.minecraft.core.Registry;
import net.minecraft.data.BuiltinRegistries;
import net.minecraft.data.worldgen.Pools;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.tags.BiomeTags;
import net.minecraft.tags.TagKey;
import net.minecraft.world.level.biome.Biome;
import net.minecraft.world.level.levelgen.GenerationStep.Decoration;
import net.minecraft.world.level.levelgen.feature.ConfiguredStructureFeature;
import net.minecraft.world.level.levelgen.feature.StructureFeature;
import net.minecraft.world.level.levelgen.feature.configurations.FeatureConfiguration;
import net.minecraft.world.level.levelgen.feature.configurations.JigsawConfiguration;
import net.minecraft.world.level.levelgen.structure.pools.LegacySinglePoolElement;
import net.minecraft.world.level.levelgen.structure.pools.StructurePoolElement;
import net.minecraft.world.level.levelgen.structure.pools.StructureTemplatePool;
import net.minecraft.world.level.levelgen.structure.pools.StructureTemplatePool.Projection;
import net.minecraftforge.common.MinecraftForge;
import net.minecraftforge.common.Tags.Biomes;
import net.minecraftforge.common.world.BiomeGenerationSettingsBuilder;
import net.minecraftforge.event.world.BiomeLoadingEvent;
import net.minecraftforge.eventbus.api.EventPriority;
import net.minecraftforge.eventbus.api.IEventBus;
import net.minecraftforge.fml.event.lifecycle.FMLCommonSetupEvent;
import net.minecraftforge.registries.DeferredRegister;
import net.minecraftforge.registries.ForgeRegistries;
import net.minecraftforge.registries.RegistryObject;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Supplier;

public class WorldGenManager {

	public static final Map<String, RegistryObject<StructureFeature<?>>> STRUCTURE_FEATURES = new HashMap<>();
	private static final DeferredRegister<StructureFeature<?>> STRUCTURE = DeferredRegister.create(ForgeRegistries.STRUCTURE_FEATURES, PVZReborn.MOD_ID);
	private static final DeferredRegister<ConfiguredStructureFeature<?, ?>> CONFIGURED_STRUCTURE_FEATURES = DeferredRegister.create(Registry.CONFIGURED_STRUCTURE_FEATURE_REGISTRY, PVZReborn.MOD_ID);

	public static void init(IEventBus bus) {
		PVZReborn.LOGGER.info("Initializing world generator");
		STRUCTURE.register(bus);
		CONFIGURED_STRUCTURE_FEATURES.register(bus);
		bus.addListener(WorldGenManager::setup);
		MinecraftForge.EVENT_BUS.addListener(EventPriority.LOWEST, WorldGenManager::onBiomeLoading);

		registerFeature("ice_tower", StructureFeatureBase::new);
		registerConfigured("ice_tower", new JigsawConfiguration(registerPool("ice_tower"), 1), Biomes.IS_SNOWY);

		registerFeature("knight_tomb", KnightTomb::new);
		registerConfigured("knight_tomb", new JigsawConfiguration(registerPool("knight_tomb"), 1), BiomeTags.HAS_VILLAGE_PLAINS);

		registerFeature("nether_tower", NetherTower::new);
		registerConfigured("nether_tower", new JigsawConfiguration(registerPool("nether_tower"), 1), BiomeTags.IS_NETHER);
		PVZReborn.LOGGER.info("Initialized world generator");
	}

	public static <C extends FeatureConfiguration> void registerFeature(String name, Supplier<? extends StructureFeature<C>> supplier) {
		WorldGenManager.STRUCTURE_FEATURES.put(name, WorldGenManager.STRUCTURE.register(name, supplier));
	}

	public static <C extends FeatureConfiguration> void registerConfigured(String name, C configured, TagKey<Biome> biome) {
		WorldGenManager.CONFIGURED_STRUCTURE_FEATURES.register(name, () -> getFeature(name).configured(configured, biome));
	}

	public static Holder<StructureTemplatePool> registerPool(String name) {
		return Pools.register(new StructureTemplatePool(PVZReborn.prefix(name), new ResourceLocation("empty"), List.of(Pair.of(StructurePoolElement.single(PVZReborn.MOD_ID_WITH_COLON + name), 1)), Projection.RIGID));
	}

	@SuppressWarnings("unchecked")
	public static <C extends FeatureConfiguration> StructureFeature<C> getFeature(String name) {
		return (StructureFeature<C>) STRUCTURE_FEATURES.get(name).get();
	}

	private static void setup(FMLCommonSetupEvent event) {
		event.enqueueWork(() -> {
			LegacySinglePoolElement masterStatues = createPoolElement("master_statues");
			LegacySinglePoolElement ironGolemStatues = createPoolElement("iron_golem_statues");
			for (String biome : new String[]{"plains", "snowy", "savanna", "desert", "taiga"}) {
				ResourceLocation poolName = new ResourceLocation("village/" + biome + "/houses");
				StructureTemplatePool pool = BuiltinRegistries.TEMPLATE_POOL.get(poolName);

				if (pool != null) {
					List<Pair<StructurePoolElement, Integer>> rawTemplates = pool.rawTemplates;
					rawTemplates.add(Pair.of(ironGolemStatues, 2));
					rawTemplates.add(Pair.of(masterStatues, 2));
					pool.templates.add(ironGolemStatues);
					pool.templates.add(ironGolemStatues);
					pool.templates.add(masterStatues);
					pool.templates.add(masterStatues);
				}
			}
			TreeFeatures.registerTreeConfig();
			Placements.registerPlacement();
		});
	}

	private static LegacySinglePoolElement createPoolElement(String resourceLocation) {
		return StructurePoolElement.legacy(PVZReborn.MOD_ID_WITH_COLON + resourceLocation).apply(Projection.RIGID);
	}

	private static void onBiomeLoading(BiomeLoadingEvent event) {
		ResourceLocation name = event.getName();
		BiomeGenerationSettingsBuilder generation = event.getGeneration();
		if (name != null && name.getNamespace().equals("minecraft")) {
			String path = name.getPath();
			if (path.equals("forest")) {
				generation.addFeature(Decoration.VEGETAL_DECORATION, Placements.NUT);
				generation.addFeature(Decoration.VEGETAL_DECORATION, Placements.CHERRY);
			} else if (path.equals("jungle")) {
				generation.addFeature(Decoration.VEGETAL_DECORATION, Placements.STAR_FRUIT);
			}
		}
	}

}
